#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
流程图识别程序主文件
按照需求文档实现从图片输入到JSON输出的完整流程
"""

import argparse
import logging
import sys
import os
from typing import Optional
import time

# 导入自定义模块
from src.preprocessing import ImageProcessor
from src.detection import NodeDetector
from src.connection import ConnectionDetector
from src.ocr import TextRecognizer
from src.analysis import LogicAnalyzer
from src.output import JSONGenerator

class FlowchartRecognizer:
    """流程图识别主程序"""
    
    def __init__(self, tesseract_path: Optional[str] = None):
        self.logger = self._setup_logging()
        
        # 初始化各个模块，为ImageProcessor传入优化参数
        self.image_processor = ImageProcessor(max_size=4000)
        self.node_detector = NodeDetector()
        self.connection_detector = ConnectionDetector()
        self.text_recognizer = TextRecognizer(tesseract_path)
        self.logic_analyzer = LogicAnalyzer()
        self.json_generator = JSONGenerator()
        
        self.logger.info("流程图识别程序初始化完成")
    
    def _setup_logging(self) -> logging.Logger:
        """设置日志配置"""
        logging.basicConfig(
            level=logging.INFO,
            format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
            handlers=[
                logging.StreamHandler(sys.stdout),
                logging.FileHandler('flowchart_recognizer.log', encoding='utf-8')
            ]
        )
        return logging.getLogger(__name__)
    
    def process_flowchart(self, image_path: str, output_dir: Optional[str] = None) -> bool:
        """处理流程图的主流程"""
        start_time = time.time()
        
        try:
            self.logger.info(f"开始处理流程图: {image_path}")
            
            # 第一步：图像预处理
            self.logger.info("=== 第一步：图像预处理 ===")
            processed_image, binary_image, metadata = self.image_processor.preprocess_image(image_path)
            
            # 检查图像质量
            if metadata['quality_score'] < 0.3:
                self.logger.warning(f"图像质量较低: {metadata['quality_score']:.2f}")
            
            # 第二步：节点检测
            self.logger.info("=== 第二步：节点检测 ===")
            nodes = self.node_detector.detect_nodes_with_filtering(processed_image, binary_image)
            
            if not nodes:
                self.logger.error("未检测到任何节点")
                return self._handle_error("未检测到任何节点", {"reason": "节点检测失败", "quality_score": metadata['quality_score']})
            
            # 第三步：连接线检测
            self.logger.info("=== 第三步：连接线检测 ===")
            connections = self.connection_detector.detect_connections_with_nodes(
                processed_image, binary_image, nodes
            )
            
            # 第四步：文本识别
            self.logger.info("=== 第四步：文本识别 ===")
            nodes_with_text = self.text_recognizer.batch_extract_text(processed_image, nodes)
            
            # 第五步：逻辑分析
            self.logger.info("=== 第五步：逻辑分析 ===")
            analysis_result = self.logic_analyzer.analyze_flowchart_structure(nodes_with_text, connections)
            
            # 第六步：生成JSON输出
            self.logger.info("=== 第六步：生成JSON输出 ===")
            json_data = self.json_generator.generate_flowchart_json(
                image_path, nodes_with_text, connections, analysis_result, metadata
            )
            
            # 验证JSON结构
            if not self.json_generator.validate_json_structure(json_data):
                self.logger.error("JSON结构验证失败")
                return self._handle_error("JSON结构验证失败", {"reason": "输出格式错误"})
            
            # 保存JSON文件
            output_path = self.json_generator.generate_output_path(image_path, output_dir)
            if not self.json_generator.save_json_to_file(json_data, output_path):
                return self._handle_error("保存JSON文件失败", {"reason": "文件写入错误"})
            
            # 计算处理时间
            processing_time = time.time() - start_time
            
            # 输出处理结果摘要
            self._print_summary(json_data, processing_time, output_path)
            
            self.logger.info(f"流程图处理完成，耗时: {processing_time:.2f}秒")
            return True
            
        except Exception as e:
            self.logger.error(f"处理流程图时发生错误: {e}")
            output_dir_for_error = output_dir if output_dir else "output"
            return self._handle_error(f"处理失败: {str(e)}", {"reason": "程序异常"}, output_dir_for_error)
    
    def _handle_error(self, error_message: str, details: dict, output_dir: str = "output") -> bool:
        """处理错误情况"""
        error_json = self.json_generator.generate_error_json(error_message, details)
        
        # 确保输出目录存在
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)

        # 保存错误JSON文件
        error_output_path = os.path.join(output_dir, "error_output.json")
        self.json_generator.save_json_to_file(error_json, error_output_path)
        
        self.logger.error(f"错误信息已保存到: {error_output_path}")
        return False
    
    def _print_summary(self, json_data: dict, processing_time: float, output_path: str):
        """打印处理结果摘要"""
        metadata = json_data['metadata']
        nodes = json_data['nodes']
        connections = json_data['connections']
        summary = json_data['structure']['summary']
        
        print("\n" + "="*50)
        print("流程图识别结果摘要")
        print("="*50)
        print(f"输入文件: {metadata['filename']}")
        print(f"输出文件: {output_path}")
        print(f"处理时间: {processing_time:.2f}秒")
        print(f"图像尺寸: {metadata['width']} x {metadata['height']}")
        print(f"整体置信度: {metadata['confidence_score']:.2f}")
        print(f"检测节点数: {len(nodes)}")
        print(f"连接线数: {len(connections)}")
        print(f"决策节点数: {summary['decision_count']}")
        print(f"流程复杂度: {summary['complexity']}")
        print(f"是否包含循环: {'是' if summary['has_cycles'] else '否'}")
        print("="*50)
        
        # 打印节点信息
        print("\n节点详情:")
        for node in nodes:
            print(f"  {node['id']}: {node['type']} - '{node['text']}' (置信度: {node['confidence']:.2f})")
        
        print("\n处理完成！")

def main():
    """主函数"""
    parser = argparse.ArgumentParser(description='流程图识别程序')
    parser.add_argument('--image_path', required=True, help='输入流程图图片路径')
    parser.add_argument('--output_dir', help='输出目录（可选）')
    parser.add_argument('--tesseract_path', help='Tesseract OCR路径（可选）')
    parser.add_argument('--verbose', '-v', action='store_true', help='详细输出模式')
    
    args = parser.parse_args()
    
    # 检查输入文件是否存在
    if not os.path.exists(args.image_path):
        print(f"错误: 输入文件不存在: {args.image_path}")
        sys.exit(1)
    
    # 检查文件格式
    supported_formats = ['.jpg', '.jpeg', '.png', '.bmp']
    file_ext = os.path.splitext(args.image_path)[1].lower()
    if file_ext not in supported_formats:
        print(f"错误: 不支持的文件格式: {file_ext}")
        print(f"支持的格式: {', '.join(supported_formats)}")
        sys.exit(1)
    
    # 设置日志级别
    if args.verbose:
        logging.getLogger().setLevel(logging.DEBUG)
    
    # 创建识别器实例
    recognizer = FlowchartRecognizer(args.tesseract_path)
    
    # 处理流程图
    success = recognizer.process_flowchart(args.image_path, args.output_dir)
    
    if success:
        print("\n✅ 流程图识别成功完成！")
        sys.exit(0)
    else:
        print("\n❌ 流程图识别失败！")
        sys.exit(1)

if __name__ == "__main__":
    main() 